﻿using TorchSharp;
using static TorchSharp.torch;

namespace Qwen3.Module;

public class Qwen3RotaryEmbedding : nn.Module<torch.Tensor, torch.Tensor, (torch.Tensor, torch.Tensor)>
{
    public Qwen3RotaryEmbedding(Qwen3Config config) : base(nameof(Qwen3RotaryEmbedding))
    {
        var dim = config.HeadDim;
        var ropeTheta = config.RopeTheta;
        var thetaNumerator = torch.arange(0, dim, 2, dtype: ScalarType.Int64).to(torch.float32);
        this.register_buffer("inv_freq", torch.pow(ropeTheta, -1.0f * (thetaNumerator / dim)), persistent: false);
    }

    public override (torch.Tensor, torch.Tensor) forward(torch.Tensor x, torch.Tensor positionIds)
    {
        using var _ = NewDisposeScope();
        Tensor invFreq;
        // Check the dimensionality of position_ids to decide shape
        // shape is typically B x T (2D) for text, or B x 3 x T (3D) for multimodal
        if (positionIds.ndim == 3)
        {
            invFreq = this.get_buffer("inv_freq").unsqueeze(0).unsqueeze(0).to(x.device);
        }
        else
        {
            invFreq = this.get_buffer("inv_freq").to(x.device);
        }
        var positionIdsExpanded = positionIds.unsqueeze(-1);
        var freqs = invFreq * positionIdsExpanded;
        var emb = torch.cat([freqs, freqs], dim: -1);
        var cos = torch.cos(emb);
        var sin = torch.sin(emb);
        return (cos.to_type(x.dtype).MoveToOuterDisposeScope(), sin.to_type(x.dtype).MoveToOuterDisposeScope());
    }
}